#!/usr/bin/env python3
"""
机器学习模型训练器模块

本模块提供了完整的机器学习模型训练和预测功能，专门用于编程问题类型的自动分类。
支持多种机器学习算法，提供模型训练、评估、保存、加载和预测的完整流程。

主要功能：
- 多种机器学习算法支持：朴素贝叶斯、逻辑回归、随机森林
- 文本特征提取：使用TF-IDF向量化技术
- 模型训练和评估：包含交叉验证和性能指标
- 模型持久化：支持模型的保存和加载
- 批量预测：支持单个和批量文本预测
- 特征重要性分析：分析关键特征对分类的影响

技术特点：
- 使用scikit-learn机器学习库
- TF-IDF文本向量化，支持n-gram特征
- 分层抽样确保训练集和测试集的类别平衡
- 交叉验证评估模型泛化能力
- 支持中文文本处理
- 模型性能自动比较和最佳模型选择

算法支持：
- 朴素贝叶斯（MultinomialNB）：适合文本分类，计算效率高
- 逻辑回归（LogisticRegression）：线性分类器，可解释性强
- 随机森林（RandomForestClassifier）：集成学习，性能稳定

使用场景：
- 编程问题自动分类
- 智能教学辅助系统
- 在线编程平台
- 编程竞赛题目管理

依赖关系：
- scikit-learn：机器学习算法库
- numpy：数值计算
- joblib：模型序列化
- pickle：数据序列化
- ml_dataset：训练数据集
"""

import os
import pickle
import joblib
import numpy as np
from typing import Dict, List, Any, Tuple, Optional
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.naive_bayes import MultinomialNB
from sklearn.linear_model import LogisticRegression
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import warnings
warnings.filterwarnings('ignore')

from ml_dataset import ProblemDataset

class MLProblemClassifier:
    """
    机器学习问题分类器
    
    提供基于机器学习的编程问题自动分类功能。支持多种算法，
    使用TF-IDF进行文本特征提取，提供完整的训练、预测和模型管理功能。
    
    主要功能：
    1. 模型训练：支持朴素贝叶斯、逻辑回归、随机森林算法
    2. 文本向量化：使用TF-IDF提取文本特征
    3. 模型评估：提供准确率、交叉验证等性能指标
    4. 预测功能：支持单个和批量文本预测
    5. 模型持久化：支持模型的保存和加载
    6. 特征分析：分析特征重要性（适用于支持的模型）
    
    技术特点：
    - TF-IDF向量化：最大5000个特征，支持1-2gram
    - 分层抽样：确保训练集和测试集的类别平衡
    - 交叉验证：评估模型泛化能力
    - 置信度计算：提供预测的可信度评估
    - 中文支持：针对中文文本优化
    
    支持的模型类型：
    - naive_bayes: 朴素贝叶斯，适合文本分类
    - logistic_regression: 逻辑回归，线性分类器
    - random_forest: 随机森林，集成学习算法
    
    使用示例：
        # 创建分类器
        classifier = MLProblemClassifier("random_forest")
        
        # 训练模型
        texts = ["计算两个数的和", "找到数组最大值"]
        labels = ["addition", "array_max"]
        results = classifier.train(texts, labels)
        
        # 预测
        prediction, confidence = classifier.predict("计算三个数的和")
        print(f"预测: {prediction}, 置信度: {confidence}")
        
        # 保存模型
        classifier.save_model("model.joblib")
    """
    
    def __init__(self, model_type: str = "random_forest"):
        """
        初始化机器学习问题分类器
        
        创建分类器实例，设置文本向量化器和模型类型。
        使用TF-IDF进行文本特征提取，支持中文文本处理。
        
        参数:
            model_type (str): 模型类型，可选值：
                - "naive_bayes": 朴素贝叶斯分类器
                - "logistic_regression": 逻辑回归分类器  
                - "random_forest": 随机森林分类器（默认）
                
        属性:
            model_type (str): 当前使用的模型类型
            vectorizer (TfidfVectorizer): TF-IDF文本向量化器
            model: 机器学习模型实例
            label_to_index (dict): 标签到索引的映射
            index_to_label (dict): 索引到标签的映射
            is_trained (bool): 模型是否已训练
            
        使用示例:
            classifier = MLProblemClassifier("random_forest")
            print(f"模型类型: {classifier.model_type}")
        """
        self.model_type = model_type
        
        # 初始化TF-IDF向量化器
        # 针对中文文本优化，不使用英文停用词
        self.vectorizer = TfidfVectorizer(
            max_features=5000,      # 最大特征数，控制特征维度
            ngram_range=(1, 2),     # 使用1-gram和2-gram特征
            stop_words=None,        # 中文不需要英文停用词
            min_df=1,               # 最小文档频率
            max_df=0.95             # 最大文档频率，过滤过于常见的词
        )
        
        # 初始化模型和标签映射
        self.model = None
        self.label_to_index = {}
        self.index_to_label = {}
        self.is_trained = False
        
    def _create_model(self):
        """
        创建指定类型的机器学习模型
        
        根据model_type参数创建相应的机器学习模型实例。
        每种模型都有针对性的参数调优。
        
        返回:
            机器学习模型实例
            
        异常:
            ValueError: 当模型类型不支持时抛出
            
        支持的模型类型:
            - naive_bayes: 朴素贝叶斯分类器，适合文本分类
            - logistic_regression: 逻辑回归分类器，线性分类器
            - random_forest: 随机森林分类器，集成学习算法
            
        使用示例:
            classifier = MLProblemClassifier("random_forest")
            model = classifier._create_model()
            print(f"创建的模型类型: {type(model)}")
        """
        if self.model_type == "naive_bayes":
            # 朴素贝叶斯分类器，alpha=0.1进行拉普拉斯平滑
            return MultinomialNB(alpha=0.1)
        elif self.model_type == "logistic_regression":
            # 逻辑回归分类器，设置随机种子和最大迭代次数
            return LogisticRegression(random_state=42, max_iter=1000)
        elif self.model_type == "random_forest":
            # 随机森林分类器，集成100棵决策树
            return RandomForestClassifier(
                n_estimators=100,        # 决策树数量
                random_state=42,         # 随机种子
                max_depth=10,            # 最大深度
                min_samples_split=5,     # 分裂所需最小样本数
                min_samples_leaf=2       # 叶节点最小样本数
            )
        else:
            raise ValueError(f"不支持的模型类型: {self.model_type}")
    
    def train(self, texts: List[str], labels: List[str], test_size: float = 0.2) -> Dict[str, Any]:
        """
        训练机器学习模型
        
        使用提供的文本和标签数据训练分类器。包含完整的训练流程：
        文本向量化、数据分割、模型训练、性能评估和交叉验证。
        
        参数:
            texts (List[str]): 问题文本列表，用于训练
            labels (List[str]): 对应的标签列表，与texts一一对应
            test_size (float, optional): 测试集比例，默认0.2（20%）
            
        返回:
            Dict[str, Any]: 训练结果字典，包含以下键：
                - model_type: 模型类型
                - train_accuracy: 训练集准确率
                - test_accuracy: 测试集准确率
                - cv_mean: 交叉验证平均准确率
                - cv_std: 交叉验证标准差
                - classification_report: 详细分类报告
                - confusion_matrix: 混淆矩阵
                - feature_count: 特征数量
                - label_count: 标签类别数
                - training_samples: 训练样本数
                - test_samples: 测试样本数
                
        异常:
            ValueError: 当输入数据格式不正确时抛出
            
        训练流程:
            1. 创建标签映射（标签到索引的转换）
            2. 文本向量化（TF-IDF特征提取）
            3. 数据分割（训练集和测试集）
            4. 模型训练
            5. 性能评估（准确率、交叉验证）
            6. 生成分类报告和混淆矩阵
            
        使用示例:
            classifier = MLProblemClassifier("random_forest")
            texts = ["计算两个数的和", "找到数组最大值", "反转字符串"]
            labels = ["addition", "array_max", "string_reverse"]
            results = classifier.train(texts, labels, test_size=0.3)
            print(f"测试准确率: {results['test_accuracy']:.4f}")
        """
        print(f"开始训练 {self.model_type} 模型...")
        print(f"训练样本数: {len(texts)}")
        
        # 创建标签映射
        # 将字符串标签转换为数字索引，便于模型处理
        unique_labels = list(set(labels))
        self.label_to_index = {label: idx for idx, label in enumerate(unique_labels)}
        self.index_to_label = {idx: label for label, idx in self.label_to_index.items()}
        
        print(f"标签类别数: {len(unique_labels)}")
        
        # 文本向量化
        # 使用TF-IDF将文本转换为数值特征向量
        print("正在进行文本向量化...")
        X = self.vectorizer.fit_transform(texts)
        y = [self.label_to_index[label] for label in labels]
        
        print(f"特征维度: {X.shape}")
        
        # 分割训练集和测试集
        # 检查是否有足够的样本进行分层抽样
        unique_labels = list(set(labels))
        min_samples_per_class = min(labels.count(label) for label in unique_labels)
        test_samples_needed = int(len(texts) * test_size)
        
        if min_samples_per_class < 2 or test_samples_needed < len(unique_labels):
            # 样本数不足，使用随机分割
            print(f"警告: 样本数不足，无法进行分层抽样，使用随机分割")
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=test_size, random_state=42
            )
        else:
            # 使用分层抽样，确保训练集和测试集的类别分布一致
            X_train, X_test, y_train, y_test = train_test_split(
                X, y, test_size=test_size, random_state=42, stratify=y
            )
        
        print(f"训练集大小: {X_train.shape[0]}")
        print(f"测试集大小: {X_test.shape[0]}")
        
        # 创建和训练模型
        self.model = self._create_model()
        print(f"正在训练 {self.model_type} 模型...")
        self.model.fit(X_train, y_train)
        
        # 评估模型性能
        train_score = self.model.score(X_train, y_train)
        test_score = self.model.score(X_test, y_test)
        
        # 交叉验证评估
        # 检查是否有足够的样本进行交叉验证
        unique_labels_train = list(set(y_train))
        min_samples_per_class_train = min(y_train.count(label) for label in unique_labels_train)
        
        if min_samples_per_class_train < 2:
            # 训练集中某些类别样本数不足，跳过交叉验证
            print(f"警告: 训练集中某些类别样本数不足，跳过交叉验证")
            cv_scores = np.array([train_score])  # 使用训练准确率作为替代
        else:
            # 进行交叉验证，折数不超过最小类别样本数
            cv_folds = min(5, min_samples_per_class_train)
            cv_scores = cross_val_score(self.model, X_train, y_train, cv=cv_folds)
        
        # 预测测试集
        y_pred = self.model.predict(X_test)
        
        # 生成分类报告
        # 将索引转换回标签名称
        y_test_labels = [self.index_to_label[idx] for idx in y_test]
        y_pred_labels = [self.index_to_label[idx] for idx in y_pred]
        
        classification_rep = classification_report(
            y_test_labels, y_pred_labels, output_dict=True, zero_division=0
        )
        
        # 标记模型已训练
        self.is_trained = True
        
        # 整理训练结果
        results = {
            'model_type': self.model_type,
            'train_accuracy': train_score,
            'test_accuracy': test_score,
            'cv_mean': cv_scores.mean(),
            'cv_std': cv_scores.std(),
            'classification_report': classification_rep,
            'confusion_matrix': confusion_matrix(y_test, y_pred).tolist(),
            'feature_count': X.shape[1],
            'label_count': len(unique_labels),
            'training_samples': len(texts),
            'test_samples': len(y_test)
        }
        
        # 输出训练结果摘要
        print(f"\n训练完成!")
        print(f"训练准确率: {train_score:.4f}")
        print(f"测试准确率: {test_score:.4f}")
        print(f"交叉验证平均准确率: {cv_scores.mean():.4f} ± {cv_scores.std():.4f}")
        
        return results
    
    def predict(self, text: str) -> Tuple[str, float]:
        """
        预测单个文本的标签
        
        使用训练好的模型对单个问题文本进行分类预测，
        返回预测的标签和置信度。
        
        参数:
            text (str): 要预测的问题文本
            
        返回:
            Tuple[str, float]: (预测标签, 置信度) 元组
                - 预测标签: 模型预测的问题类型
                - 置信度: 预测的可信度，范围0-1
                
        异常:
            ValueError: 当模型尚未训练时抛出
            
        预测流程:
            1. 检查模型是否已训练
            2. 将文本转换为特征向量
            3. 使用模型进行预测
            4. 计算预测置信度
            5. 返回预测结果
            
        使用示例:
            classifier = MLProblemClassifier("random_forest")
            # 假设模型已训练
            prediction, confidence = classifier.predict("计算两个数的和")
            print(f"预测: {prediction}, 置信度: {confidence:.4f}")
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练，请先调用train方法")
        
        # 文本向量化
        # 将输入文本转换为与训练时相同的特征向量
        X = self.vectorizer.transform([text])
        
        # 模型预测
        # 获取预测的类别索引
        prediction_idx = self.model.predict(X)[0]
        prediction_label = self.index_to_label[prediction_idx]
        
        # 计算置信度
        # 如果模型支持概率预测，计算预测概率
        if hasattr(self.model, 'predict_proba'):
            probabilities = self.model.predict_proba(X)[0]
            confidence = probabilities[prediction_idx]
        else:
            # 对于不支持概率预测的模型，使用默认置信度
            confidence = 0.5
        
        return prediction_label, confidence
    
    def predict_batch(self, texts: List[str]) -> List[Tuple[str, float]]:
        """
        批量预测多个文本的标签
        
        使用训练好的模型对多个问题文本进行批量分类预测，
        比单个预测更高效，适合处理大量文本。
        
        参数:
            texts (List[str]): 要预测的问题文本列表
            
        返回:
            List[Tuple[str, float]]: 预测结果列表
                每个元素为(预测标签, 置信度)元组
                
        异常:
            ValueError: 当模型尚未训练时抛出
            
        批量预测优势:
            - 更高的计算效率
            - 减少重复的向量化开销
            - 适合处理大量数据
            
        使用示例:
            classifier = MLProblemClassifier("random_forest")
            # 假设模型已训练
            texts = ["计算两个数的和", "找到数组最大值", "反转字符串"]
            results = classifier.predict_batch(texts)
            for i, (prediction, confidence) in enumerate(results):
                print(f"文本{i+1}: {texts[i]} -> {prediction} (置信度: {confidence:.4f})")
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练，请先调用train方法")
        
        # 批量文本向量化
        # 一次性处理所有文本，提高效率
        X = self.vectorizer.transform(texts)
        
        # 批量预测
        # 获取所有文本的预测类别索引
        predictions_idx = self.model.predict(X)
        predictions_labels = [self.index_to_label[idx] for idx in predictions_idx]
        
        # 批量计算置信度
        # 如果模型支持概率预测，批量计算所有预测的概率
        if hasattr(self.model, 'predict_proba'):
            probabilities = self.model.predict_proba(X)
            confidences = [prob[idx] for prob, idx in zip(probabilities, predictions_idx)]
        else:
            # 对于不支持概率预测的模型，使用默认置信度
            confidences = [0.5] * len(texts)
        
        # 组合预测结果和置信度
        return list(zip(predictions_labels, confidences))
    
    def save_model(self, filepath: str):
        """
        保存训练好的模型到文件
        
        将训练好的模型、向量化器、标签映射等所有必要信息
        保存到文件中，以便后续加载和使用。
        
        参数:
            filepath (str): 保存文件的路径，建议使用.joblib扩展名
            
        异常:
            ValueError: 当模型尚未训练时抛出
            IOError: 当文件保存失败时抛出
            
        保存内容:
            - model: 训练好的机器学习模型
            - vectorizer: TF-IDF向量化器
            - label_to_index: 标签到索引的映射
            - index_to_label: 索引到标签的映射
            - model_type: 模型类型
            - is_trained: 训练状态标志
            
        使用示例:
            classifier = MLProblemClassifier("random_forest")
            # 假设模型已训练
            classifier.save_model("my_model.joblib")
            print("模型已保存")
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练，无法保存")
        
        # 准备模型数据
        # 包含所有必要的组件，确保模型可以完整恢复
        model_data = {
            'model': self.model,                    # 训练好的模型
            'vectorizer': self.vectorizer,          # 文本向量化器
            'label_to_index': self.label_to_index,  # 标签映射
            'index_to_label': self.index_to_label,  # 反向标签映射
            'model_type': self.model_type,          # 模型类型
            'is_trained': self.is_trained           # 训练状态
        }
        
        # 使用joblib保存模型
        # joblib比pickle更适合保存scikit-learn模型
        joblib.dump(model_data, filepath)
        print(f"模型已保存到: {filepath}")
    
    def load_model(self, filepath: str):
        """
        从文件加载训练好的模型
        
        从保存的文件中恢复模型、向量化器、标签映射等所有信息，
        使分类器可以立即用于预测。
        
        参数:
            filepath (str): 模型文件的路径
            
        异常:
            FileNotFoundError: 当模型文件不存在时抛出
            IOError: 当文件加载失败时抛出
            
        加载内容:
            - model: 训练好的机器学习模型
            - vectorizer: TF-IDF向量化器
            - label_to_index: 标签到索引的映射
            - index_to_label: 索引到标签的映射
            - model_type: 模型类型
            - is_trained: 训练状态标志
            
        使用示例:
            classifier = MLProblemClassifier()
            classifier.load_model("my_model.joblib")
            prediction, confidence = classifier.predict("计算两个数的和")
            print(f"预测: {prediction}")
        """
        if not os.path.exists(filepath):
            raise FileNotFoundError(f"模型文件不存在: {filepath}")
        
        # 从文件加载模型数据
        model_data = joblib.load(filepath)
        
        # 恢复所有组件
        self.model = model_data['model']
        self.vectorizer = model_data['vectorizer']
        self.label_to_index = model_data['label_to_index']
        self.index_to_label = model_data['index_to_label']
        self.model_type = model_data['model_type']
        self.is_trained = model_data['is_trained']
        
        print(f"模型已从 {filepath} 加载")
    
    def get_feature_importance(self, top_n: int = 20) -> List[Tuple[str, float]]:
        """
        获取特征重要性分析
        
        分析哪些特征对分类结果影响最大，仅适用于支持特征重要性的模型
        （如随机森林）。帮助理解模型的决策过程和特征的重要性。
        
        参数:
            top_n (int, optional): 返回前N个最重要的特征，默认20
            
        返回:
            List[Tuple[str, float]]: 特征重要性列表
                每个元素为(特征名称, 重要性分数)元组，按重要性降序排列
                
        异常:
            ValueError: 当模型尚未训练或模型不支持特征重要性时抛出
            
        适用模型:
            - 随机森林：支持特征重要性分析
            - 朴素贝叶斯：不支持特征重要性
            - 逻辑回归：不支持特征重要性
            
        使用示例:
            classifier = MLProblemClassifier("random_forest")
            # 假设模型已训练
            importance = classifier.get_feature_importance(top_n=10)
            print("最重要的特征:")
            for feature, score in importance:
                print(f"  {feature}: {score:.4f}")
        """
        if not self.is_trained:
            raise ValueError("模型尚未训练")
        
        if not hasattr(self.model, 'feature_importances_'):
            raise ValueError(f"{self.model_type} 模型不支持特征重要性")
        
        # 获取特征名称
        # 从向量化器中获取所有特征的名称
        feature_names = self.vectorizer.get_feature_names_out()
        
        # 获取特征重要性分数
        # 从训练好的模型中获取每个特征的重要性
        importances = self.model.feature_importances_
        
        # 组合特征名称和重要性分数
        feature_importance = list(zip(feature_names, importances))
        
        # 按重要性降序排序
        feature_importance.sort(key=lambda x: x[1], reverse=True)
        
        # 返回前top_n个最重要的特征
        return feature_importance[:top_n]

class MLModelTrainer:
    """
    机器学习模型训练器
    
    提供完整的机器学习模型训练和管理功能，支持多种算法的
    自动训练、性能比较和最佳模型选择。简化了模型训练的复杂流程。
    
    主要功能：
    1. 多模型训练：自动训练朴素贝叶斯、逻辑回归、随机森林三种模型
    2. 性能比较：自动比较不同模型的性能指标
    3. 最佳模型选择：根据测试准确率自动选择最佳模型
    4. 模型保存：保存最佳模型和所有训练结果
    5. 模型加载：加载保存的最佳模型
    
    技术特点：
    - 自动化训练流程：一键训练所有支持的模型类型
    - 性能评估：提供详细的性能比较报告
    - 最佳模型选择：基于测试准确率自动选择最优模型
    - 结果持久化：保存训练结果和模型文件
    - 易于使用：简化的接口，适合快速原型开发
    
    使用场景：
    - 模型选择和比较
    - 快速原型开发
    - 自动化模型训练
    - 性能基准测试
    
    使用示例：
        # 创建训练器
        trainer = MLModelTrainer()
        
        # 训练所有模型
        results = trainer.train_all_models()
        
        # 保存最佳模型
        model_path = trainer.save_best_model()
        
        # 加载最佳模型进行预测
        best_model = trainer.load_best_model(model_path)
    """
    
    def __init__(self):
        """
        初始化机器学习模型训练器
        
        创建训练器实例，加载数据集并初始化模型和结果存储。
        
        属性:
            dataset (ProblemDataset): 训练数据集实例
            models (dict): 存储训练好的模型，键为模型类型
            results (dict): 存储训练结果，键为模型类型
            
        使用示例:
            trainer = MLModelTrainer()
            print(f"数据集样本数: {len(trainer.dataset.training_data)}")
        """
        # 加载训练数据集
        self.dataset = ProblemDataset()
        
        # 初始化存储字典
        self.models = {}      # 存储训练好的模型
        self.results = {}     # 存储训练结果
    
    def train_all_models(self, test_size: float = 0.2) -> Dict[str, Any]:
        """
        训练所有支持的模型类型
        
        自动训练朴素贝叶斯、逻辑回归、随机森林三种模型，
        比较它们的性能并选择最佳模型。
        
        参数:
            test_size (float, optional): 测试集比例，默认0.2（20%）
            
        返回:
            Dict[str, Any]: 所有模型的训练结果字典
                键为模型类型，值为对应的训练结果
                
        训练流程:
            1. 获取训练数据（文本和标签）
            2. 依次训练三种模型类型
            3. 比较模型性能
            4. 选择最佳模型
            5. 输出性能报告
            
        使用示例:
            trainer = MLModelTrainer()
            results = trainer.train_all_models(test_size=0.3)
            
            # 查看各模型性能
            for model_type, result in results.items():
                print(f"{model_type}: {result['test_accuracy']:.4f}")
        """
        print("开始训练所有模型...")
        
        # 获取训练数据
        # 从数据集中提取文本和标签
        texts, labels = self.dataset.get_texts_and_labels()
        
        # 定义要训练的模型类型
        # 按性能从低到高的顺序排列，便于比较
        model_types = ["naive_bayes", "logistic_regression", "random_forest"]
        
        # 依次训练每种模型
        for model_type in model_types:
            print(f"\n{'='*50}")
            print(f"训练 {model_type} 模型")
            print(f"{'='*50}")
            
            # 创建分类器并训练
            classifier = MLProblemClassifier(model_type)
            results = classifier.train(texts, labels, test_size)
            
            # 保存模型和结果
            self.models[model_type] = classifier
            self.results[model_type] = results
        
        # 比较模型性能
        print(f"\n{'='*50}")
        print("模型性能比较")
        print(f"{'='*50}")
        
        # 输出各模型的性能指标
        for model_type, results in self.results.items():
            print(f"{model_type}:")
            print(f"  测试准确率: {results['test_accuracy']:.4f}")
            print(f"  交叉验证: {results['cv_mean']:.4f} ± {results['cv_std']:.4f}")
        
        # 选择最佳模型
        # 基于测试准确率选择性能最好的模型
        best_model = max(self.results.items(), key=lambda x: x[1]['test_accuracy'])
        print(f"\n最佳模型: {best_model[0]} (准确率: {best_model[1]['test_accuracy']:.4f})")
        
        return self.results
    
    def save_best_model(self, model_dir: str = "models"):
        """
        保存最佳模型和训练结果
        
        根据测试准确率选择最佳模型，将其保存到指定目录，
        同时保存所有模型的训练结果供后续分析。
        
        参数:
            model_dir (str, optional): 模型保存目录，默认"models"
            
        返回:
            str: 最佳模型的保存路径
            
        异常:
            ValueError: 当没有训练好的模型时抛出
            IOError: 当文件保存失败时抛出
            
        保存内容:
            - 最佳模型文件：best_model_{model_type}.joblib
            - 训练结果文件：training_results.json
            
        使用示例:
            trainer = MLModelTrainer()
            trainer.train_all_models()
            model_path = trainer.save_best_model("my_models")
            print(f"最佳模型保存在: {model_path}")
        """
        if not self.models:
            raise ValueError("没有训练好的模型")
        
        # 创建模型目录
        # 确保目录存在，如果不存在则创建
        os.makedirs(model_dir, exist_ok=True)
        
        # 选择最佳模型
        # 基于测试准确率选择性能最好的模型
        best_model_type = max(self.results.items(), key=lambda x: x[1]['test_accuracy'])[0]
        best_model = self.models[best_model_type]
        
        # 保存最佳模型
        # 使用模型类型作为文件名的一部分
        model_path = os.path.join(model_dir, f"best_model_{best_model_type}.joblib")
        best_model.save_model(model_path)
        
        # 保存所有模型结果
        # 将训练结果保存为JSON文件，便于后续分析
        results_path = os.path.join(model_dir, "training_results.json")
        import json
        with open(results_path, 'w', encoding='utf-8') as f:
            json.dump(self.results, f, ensure_ascii=False, indent=2)
        
        print(f"最佳模型已保存到: {model_path}")
        print(f"训练结果已保存到: {results_path}")
        
        return model_path
    
    def load_best_model(self, model_path: str) -> MLProblemClassifier:
        """
        加载保存的最佳模型
        
        从指定路径加载之前保存的最佳模型，返回可用的分类器实例。
        
        参数:
            model_path (str): 模型文件的路径
            
        返回:
            MLProblemClassifier: 加载的分类器实例，可直接用于预测
            
        异常:
            FileNotFoundError: 当模型文件不存在时抛出
            IOError: 当文件加载失败时抛出
            
        使用示例:
            trainer = MLModelTrainer()
            # 假设已有保存的模型
            best_model = trainer.load_best_model("models/best_model_random_forest.joblib")
            prediction, confidence = best_model.predict("计算两个数的和")
            print(f"预测: {prediction}")
        """
        # 创建新的分类器实例
        classifier = MLProblemClassifier()
        
        # 加载模型
        classifier.load_model(model_path)
        
        return classifier

if __name__ == "__main__":
    """
    主程序入口
    
    当直接运行此文件时，会执行完整的机器学习模型训练和测试流程：
    1. 创建模型训练器
    2. 训练所有支持的模型类型
    3. 保存最佳模型
    4. 加载最佳模型进行测试
    5. 使用测试用例验证模型性能
    
    这个演示程序展示了机器学习训练器的完整使用流程。
    """
    
    # 创建模型训练器实例
    trainer = MLModelTrainer()
    
    # 训练所有模型
    # 自动训练朴素贝叶斯、逻辑回归、随机森林三种模型
    print("开始训练所有模型...")
    results = trainer.train_all_models()
    
    # 保存最佳模型
    # 根据测试准确率选择最佳模型并保存
    print("\n保存最佳模型...")
    model_path = trainer.save_best_model()
    
    # 测试最佳模型
    print(f"\n{'='*50}")
    print("测试最佳模型")
    print(f"{'='*50}")
    
    # 加载最佳模型
    best_model = trainer.load_best_model(model_path)
    
    # 定义测试用例
    # 涵盖不同类型的编程问题，验证模型的分类能力
    test_cases = [
        "计算两个整数的和",        # 数学运算 - addition
        "找到数组中的最大值",      # 数组操作 - array_max
        "反转字符串",             # 字符串操作 - string_reverse
        "计算阶乘",               # 数学问题 - factorial
        "判断一个数是否为质数"     # 数学问题 - prime_check
    ]
    
    # 测试每个用例
    print("测试用例预测结果:")
    for i, test_case in enumerate(test_cases, 1):
        # 使用最佳模型进行预测
        prediction, confidence = best_model.predict(test_case)
        
        # 输出预测结果
        print(f"{i}. 问题: {test_case}")
        print(f"   预测: {prediction} (置信度: {confidence:.4f})")
        print()
